# -*- coding: utf-8 -*-

from conf.logger import fort_logger
from channels.generic.websocket import WebsocketConsumer
from Exscript import Account
from Exscript.protocols import Telnet
from utils.ssh import MyThread


class MyTelnet(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super(MyTelnet, self).__init__(*args, **kwargs)
        self.telnet = Telnet()
        self.ip = None
        self.port = None
        self.username = None
        self.password = None
        self.record_dir = 'admin_telnet_records' if self.scope['user'].is_superuser else 'fort_telnet_records'
        self.remote_ip = self.scope['query_string'].decode('utf-8')
        self.chan = None
        self.cmd = []  # 所有命令
        self.cmd_tmp = ''  # 一行命令
        self.tab_mode = False  # 使用tab命令补全时需要读取返回数据然后添加到当前输入命令后
        self.history_mode = False
        self.index = 0
        self.coding = 'utf-8'
        self.title = 'telnet'
        self.width = 150
        self.height = 30
        self.t1 = MyThread(self)

    def connect(self):
        if self.scope["user"].is_anonymous:
            self.close(code=1007)
        else:
            self.accept()

        try:
            self.telnet.connect(self.ip, int(self.port))
            self.telnet.login(Account(self.username, self.password))
        except Exception as e:
            fort_logger.error('用户{}通过webtelnet连接{}失败！原因：{}'.format(self.username, self.ip, e))
            self.send('用户{}通过webtelnet连接{}失败！原因：{}'.format(self.username, self.ip, e), close=True)
        self.chan = self.telnet.tn.sock
        # 设置如果3分钟没有任何输入，就断开连接
        self.chan.settimeout(60 * 3)
        self.t1.setDaemon(True)
        self.t1.start()

    def receive(self, text_data=None, bytes_data=None):
        self.chan.send(text_data.encode(encoding=self.coding))
        self.gen_cmd(text_data)

    def disconnect(self, close_code):
        try:
            self.handle_cmd()
            self.t1.record()
        finally:
            self.telnet.close()
            self.t1.stop()

    def gen_cmd(self, text_data):
        if text_data == '\r':
            self.index = 0
            if self.cmd_tmp.strip() != '':
                self.cmd.append(self.cmd_tmp)
                self.cmd_tmp = ''
        elif text_data.encode(encoding=self.coding) == b'\x07':
            pass
        elif text_data.encode(encoding=self.coding) in (b'\x03', b'\x01'):  # ctrl+c 和 ctrl+a
            self.index = 0
        elif text_data.encode(encoding=self.coding) == b'\x05':  # ctrl+e
            self.index = len(self.cmd_tmp) - 2
        elif text_data.encode(encoding=self.coding) == b'\x1b[D':  # ← 键
            if self.index == 0:
                self.index = len(self.cmd_tmp) - 2
            else:
                self.index -= 1
        elif text_data.encode(encoding=self.coding) == b'\x1b[C':  # → 键
            self.index += 1
        elif text_data.encode(encoding=self.coding) == b'\x7f':  # Backspace键
            if self.index == 0:
                self.cmd_tmp = self.cmd_tmp[:-1]
            else:
                self.cmd_tmp = self.cmd_tmp[:self.index] + self.cmd_tmp[self.index + 1:]
        else:
            if text_data == '\t' or text_data.encode(encoding=self.coding) == b'\x1b':  # \x1b 点击2下esc键也可以补全
                self.tab_mode = True
            elif text_data.encode(encoding=self.coding) == b'\x1b[A' or text_data.encode(encoding=self.coding) == b'\x1b[B':
                self.history_mode = True
            else:
                if self.index == 0:
                    self.cmd_tmp += text_data
                else:
                    self.cmd_tmp = self.cmd_tmp[:self.index] + text_data + self.cmd_tmp[self.index:]

    def handle_cmd(self):  # 将vim或vi编辑文档时的操作去掉
        vi_index = None
        fg_index = None  # 捕捉使用ctrl+z将vim放到后台的操作
        q_index = None
        q_keys = (':wq', ':q', ':q!')
        for index, value in enumerate(self.cmd):
            if 'vi' in value:
                vi_index = index
            if any([key in value for key in q_keys]):
                q_index = index
            if '\x1a' in value:  # \x1a代表ctrl+z
                self.cmd[index] = value.split('\x1a')[1]
            if 'fg' in value:
                fg_index = index

        first_index = fg_index if fg_index else vi_index
        if vi_index:
            self.cmd = self.cmd[:first_index + 1] + self.cmd[q_index + 1:]
